from docx import Document
from docx.shared import Pt, RGBColor
from docx.enum.text import WD_PARAGRAPH_ALIGNMENT
from docx.enum.table import WD_TABLE_ALIGNMENT
from utils.analysis.mediation.mediation_model6 import MediationModel6, Model6Result
from utils.report.base import format_table, create_table
import pandas as pd


def _generate_report(data, document, table_count):
    table1 = generate_table1(data, document, table_count)  # 表1
    table2 = generate_table2(data, document, table_count)  # 表2
    # 总结段落
    p2 = ("")
    document.add_paragraph(p2)  # 插入段落
    return


def generate_table2(data, document, table_count):
    document.add_paragraph('')  # 插入段落
    # 新建表格
    init_rows = 5
    init_cols = 5
    # table = document.add_table(rows=init_rows, cols=init_cols)
    table = create_table(document, init_rows, init_cols)
    # 1. 填充1-2行
    # 填充第一行表头
    row = table.rows[0]  # 表格
    f_cell = None
    for i in range(init_cols):  # 合并单元格
        if i == 0:
            f_cell = table.cell(0, 0)
        else:
            f_cell.merge(table.cell(0, i))
    frequency_table_name = ""
    row.cells[0].text = frequency_table_name + '间接效应分析'
    table.cell(1, 0).text = "中介类型"
    table.cell(2, 0).text = "总效应"
    table.cell(3, 0).text = "直接效应"
    table.cell(4, 0).text = "间接效应"
    table.cell(1, 1).text = "Effect"
    table.cell(1, 2).text = "Boot SE"
    table.cell(1, 3).text = "BootLLCI"
    table.cell(1, 4).text = "BootULCI"
    tables = data.mediation_table
    for i in tables:
        if 'Total' == i.path:
            cur_row = 2
            table.cell(cur_row, 1).text = convert_str_to_float_3(i.coef)
            table.cell(cur_row, 2).text = convert_str_to_float_3(i.se)
            table.cell(cur_row, 3).text = convert_str_to_float_3(i.CI_2_5)
            table.cell(cur_row, 4).text = convert_str_to_float_3(i.CI_97_5)
        if 'Direct' == i.path:
            cur_row = 3
            table.cell(cur_row, 1).text = convert_str_to_float_3(i.coef)
            table.cell(cur_row, 2).text = convert_str_to_float_3(i.se)
            table.cell(cur_row, 3).text = convert_str_to_float_3(i.CI_2_5)
            table.cell(cur_row, 4).text = convert_str_to_float_3(i.CI_97_5)
        if 'Indirect' == i.path:
            cur_row = 4
            table.cell(cur_row, 1).text = convert_str_to_float_3(i.coef)
            table.cell(cur_row, 2).text = convert_str_to_float_3(i.se)
            table.cell(cur_row, 3).text = convert_str_to_float_3(i.CI_2_5)
            table.cell(cur_row, 4).text = convert_str_to_float_3(i.CI_97_5)
    # 格式化表格
    format_table(table)
    return table


def generate_table1(data, document, table_count):
    # 新建表格
    init_rows = 3
    init_cols = 9
    # table = document.add_table(rows=init_rows, cols=init_cols)
    table = create_table(document, init_rows, init_cols)
    # 1. 填充1-3行
    # 填充第一行表头
    row = table.rows[0]  # 表格
    f_cell = None
    for i in range(init_cols):  # 合并单元格
        if i == 0:
            f_cell = table.cell(0, 0)
        else:
            f_cell.merge(table.cell(0, i))
    frequency_table_name = ""
    row.cells[0].text = frequency_table_name + '中介效应模型汇总'
    table.cell(1, 0).merge(table.cell(2, 0))
    table.cell(1, 0).text = "变量"
    # 返回结果：1. M1 <- X；2. M2 <- X + M1; 3. Y <- X; 4. Y <- X + M1 + M2
    # m1 <- x
    table.cell(1, 1).merge(table.cell(1, 2))
    table.cell(1, 1).text = "模型1;" + data.m1
    table.cell(2, 1).text = "β"
    table.cell(2, 2).text = "t"
    # m2 <- x + m1
    table.cell(1, 3).merge(table.cell(1, 4))
    table.cell(1, 3).text = "模型2;" + data.m2
    table.cell(2, 3).text = "β"
    table.cell(2, 4).text = "t"
    # y <- x
    table.cell(1, 5).merge(table.cell(1, 6))
    table.cell(1, 5).text = "模型3;" + data.y
    table.cell(2, 5).text = "β"
    table.cell(2, 6).text = "t"
    # y <- x + m1 + m2
    table.cell(1, 7).merge(table.cell(1, 6))
    table.cell(1, 7).text = "模型4;" + data.y
    table.cell(2, 7).text = "β"
    table.cell(2, 8).text = "t"
    # 2. 填充4-7行
    add_values(table, "Intercept", data.m1_x.details, data.m2_x_m1.details, data.y_x.details, data.y_x_m1_m2.details)
    add_values(table, data.x, data.m1_x.details, data.m2_x_m1.details, data.y_x.details, data.y_x_m1_m2.details)
    add_values(table, data.m1, data.m1_x.details, data.m2_x_m1.details, data.y_x.details, data.y_x_m1_m2.details)
    add_values(table, data.m2, data.m1_x.details, data.m2_x_m1.details, data.y_x.details, data.y_x_m1_m2.details)
    # 3.填充控制变量
    covars = data.covar
    for i in covars:
        add_values(table, i, data.m1_x.details, data.m2_x_m1.details, data.y_x.details, data.y_x_m1_m2.details)
    # 4.填充R方，调整后的r方和F值
    # add_last_3_cols(table, 'R2', convert_str_to_float_3(m2x_r2), convert_str_to_float_3(y2x_r2),
    #                 convert_str_to_float_3(y2xm_r2))
    # add_last_3_cols(table, '调整R2', convert_str_to_float_3(m2x_adj_r2), convert_str_to_float_3(y2x_adj_r2),
    #                 convert_str_to_float_3(y2xm_adj_r2))
    # add_last_3_cols(table, 'F值', convert_str_to_float_3(data.r_m2x.f) + add_p_value(data.r_m2x.p),
    #                 convert_str_to_float_3(data.r_y2x.f) + add_p_value(data.r_y2x.p),
    #                 convert_str_to_float_3(data.r_y2xm.f) + add_p_value(data.r_y2xm.p),
    #                 )
    # 格式化表格
    format_table(table)

    table.add_row()  # 表格动态增加一行
    cur_row = table.rows[-1]
    cur_row.cells[0].text = '* p<0.05 ** p<0.01 *** p<0.001'
    cur_row.cells[0].merge(cur_row.cells[1]).merge(cur_row.cells[2]).merge(cur_row.cells[3]).merge(
        cur_row.cells[4]).merge(cur_row.cells[5]).merge(cur_row.cells[6])
    return table


def add_values(table, cur_name, details_m1, details_m2, details_m3, details_m4):
    table.add_row()  # 表格动态增加一行
    cur_row = table.rows[-1]
    cur_row.cells[0].text = cur_name
    for i in details_m1:
        if cur_name == i.names:
            cur_str = convert_str_to_float_3(i.coef) + add_p_value(i.pval)
            cur_row.cells[1].text = cur_str
            cur_row.cells[2].text = convert_str_to_float_3(i.T)
    for i in details_m2:
        if cur_name == i.names:
            cur_str = convert_str_to_float_3(i.coef) + add_p_value(i.pval)
            cur_row.cells[3].text = cur_str
            cur_row.cells[4].text = convert_str_to_float_3(i.T)
    for i in details_m3:
        if cur_name == i.names:
            cur_str = convert_str_to_float_3(i.coef) + add_p_value(i.pval)
            cur_row.cells[5].text = cur_str
            cur_row.cells[6].text = convert_str_to_float_3(i.T)
    for i in details_m4:
        if cur_name == i.names:
            cur_str = convert_str_to_float_3(i.coef) + add_p_value(i.pval)
            cur_row.cells[7].text = cur_str
            cur_row.cells[8].text = convert_str_to_float_3(i.T)


def add_last_3_cols(table, name, v1, v2, v3):
    table.add_row()  # 表格动态增加一行
    cur_row = table.rows[-1]
    cur_row.cells[0].text = name
    cur_row.cells[1].merge(cur_row.cells[2]).text = v1
    cur_row.cells[3].merge(cur_row.cells[4]).text = v2
    cur_row.cells[5].merge(cur_row.cells[6]).text = v3


def convert_str_to_float_3(value) -> str:
    if not value:
        return ''
    if isinstance(value, str):
        return "{:.3f}".format(float(value))
    elif isinstance(value, float):
        return "{:.3f}".format(value)
    return "{:.3f}".format(value)


def add_p_value(p):
    res = ''
    if p <= 0.001:
        res = '***'
    elif p <= 0.01:
        res = '**'
    elif p <= 0.05:
        res = '*'
    return res


def generate(src_data: [Model6Result], doc: Document = None) -> Document():
    if not doc:
        doc = Document()
    title1 = doc.add_heading(level=1)  # 增加标题
    t1_run = title1.add_run('中介效应分析-模型6')
    t1_run.font.color.rgb = RGBColor(10, 10, 10)
    for i in src_data:
        _generate_report(i, doc, 1)  # 生成表格
    return doc


if __name__ == '__main__':
    # 创建 Document 对象，等价于在电脑上打开一个 Word 文档
    doc = Document()
    pd.set_option('expand_frame_repr', False)
    df = pd.read_excel('./test_datas.xlsx')
    df = df[["性别", "年龄", "学历", "企业角色", 'JG', 'YY', 'JX', 'XW']]
    # 1.此文件的结果不包含整体的F值和p值
    # 2.若有多个x,y,m，则需要遍历进行获取结果
    x_arr = ['JG']
    y_arr = ['YY']
    m_arr = ['XW', 'JX']
    covar = ["性别", "年龄", "学历", "企业角色"]
    obj = MediationModel6(df, x_arr, y_arr, m_arr, covar)
    src_data = obj.analysis()
    title1 = doc.add_heading(level=1)  # 增加标题
    t1_run = title1.add_run('中介效应分析-模型6')
    p1 = ("如表所示，通过Process 宏程序中的 Bootstrap 方法, 结合逐步检验法，运用model 4进行中介作用分析，针对“M”这一中介变量，"
          "c表示X对Y时的回归系数（模型中没有中介变量M时），即总效应，当总效应显著时，讨论中介效应影响机制；a表示X对M时的回归系数，"
          "b表示M对Y时的回归系数，a*b为a与b的乘积即中介效应；c’表示X对Y时的回归系数（模型中有中介变量时），即直接效应；95% BootCI表示"
          "Bootstrap抽样计算得到的95%置信区间。如果a和b显著，c’不显著，且a*b的95% BootCI包括数字不包括0，则为完全中介；后续开始与"
          "文章背景开始结合；如果a和b显著，c’显著，且a*b的95% BootCI包括数字不包括0，则为部分中介；如果a和b至少一个不显著，"
          "且a*b的95% BootCI包括数字0（不显著），则中介作用不显著；")
    doc.add_paragraph(p1)  # 插入段落
    t1_run.font.color.rgb = RGBColor(10, 10, 10)
    for i in src_data:
        _generate_report(i, doc, 1)  # 生成表格
    # 保存文档
    doc.save('demo_zj_model6.docx')
